# 从0开始学习代码   李康

# os包: 操作文件和目录的模块
import os
# sys包: 系统相关的模块
import sys
# cv2包: OpenCV图像处理包
# numpy包: 数组计算包
# torch包: PyTorch深度学习框架
import torch
# tqdm包: 进度条
from tqdm import tqdm
#  训练参数配置文件
from config.config import Config as cfg
# 数据集读取和预处理相关包  我这里分开import是因为我用了showComment来显示注释
from data.augmentation import RandomBlur
from data.augmentation import RandomColorJitter
from data.augmentation import augCompose
from data.dataset import dataReadPip
from data.dataset import loadedDataset
from data.dataset import readIndex

from model.deepcrack_with_uu_Sp_3 import DeepCrack_uu_Sp_3
from trainer.trainer import DeepCrackTrainer


os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu_id # 设置GPU编号

"""
DeepCrack裂缝检测模型训练主程序

核心功能：
1. 数据加载与预处理
2. 多GPU模型训练
3. 训练过程可视化监控
4. 模型验证与保存
5. 异常中断处理

技术特性：
- 支持多尺度特征融合训练
- 实时可视化训练指标（损失/准确率/样本）
- 基于最佳验证指标的模型保存策略
- 键盘中断安全处理
"""

def main():
    """主训练函数"""

    # ==================== 数据准备阶段 ==================== #
    # 数据增强配置（50%概率色彩抖动+20%概率模糊）
    # data_augment_op = augCompose(transforms=[[RandomColorJitter, 0.5], [RandomBlur, 0.2]])

    # 数据读取管道（训练集/验证集）
    train_pipline = dataReadPip(transforms=None)  # 训练集不使用额外变换
    test_pipline = dataReadPip(transforms=None)  # 验证集保持原始数据

    # 数据集加载（自动打乱训练集）
    train_dataset = loadedDataset(readIndex(cfg.train_data_path, shuffle=True), preprocess=train_pipline)
    test_dataset = loadedDataset(readIndex(cfg.test_data_path), preprocess=test_pipline)

    # DataLoader配置（4线程加载，丢弃不完整batch）
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=cfg.train_batch_size,
        shuffle=True,
        num_workers=4,
        drop_last=True
    )
    val_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=cfg.val_batch_size,
        shuffle=False,
        num_workers=4,
        drop_last=True
    )

    # ==================== 模型构建阶段 ==================== #
    # 初始化DeepCrack模型（多GPU并行）
    model = DeepCrack_uu_Sp_3().cuda()
    model = torch.nn.DataParallel(model, device_ids=[0]).cuda()

    # 构建训练器（集成优化器/损失函数）
    trainer = DeepCrackTrainer(model).cuda()

    # 加载预训练模型（选择性参数更新）   若没有 预训练模型，则从0开始训练
    if cfg.pretrained_model:
        #
        pretrained_dict = trainer.saver.load(cfg.pretrained_model, multi_gpu=True)
        model_dict = model.state_dict()

        # 过滤不匹配的层参数
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)

    # ==================== 开启visdom本地日志保存 edit by likang==================== #
        trainer.vis.save_settings(save_path=cfg.vis_save_path, save_log=True, save_img=True, save_plot=True )
        trainer.vis.log('load checkpoint: %s' % cfg.pretrained_model, 'train info')

    # ==================== 训练循环阶段 ==================== #
    try:
        for epoch in range(1, cfg.epoch):
            trainer.vis.save_settings(save_path=cfg.vis_save_path, save_log=True, save_img=True, save_plot=True)
            trainer.vis.log('Start Epoch %d ...' % epoch, 'train info')
            model.train()

            # --------------------- 训练批次处理 ------------------- #
            bar = tqdm(enumerate(train_loader), total=len(train_loader))
            bar.set_description('Epoch %d --- Training --- :' % epoch)

            for idx, (img, lab) in bar:
                # 数据转移到GPU
                data = img.type(torch.cuda.FloatTensor).cuda()
                target = lab.type(torch.cuda.FloatTensor).cuda()

                # 前向传播与参数更新
                pred = trainer.train_op(data, target)

                # 定期记录训练损失（默认每40步）
                if idx % cfg.vis_train_loss_every == 0:
                    trainer.vis.log(trainer.log_loss, 'train_loss')
                    trainer.vis.plot_many({
                        'train_total_loss': trainer.log_loss['total_loss'],
                        'train_output_loss': trainer.log_loss['output_loss'],
                        'train_fuse5_loss': trainer.log_loss['fuse5_loss'],
                        'train_fuse4_loss': trainer.log_loss['fuse4_loss'],
                        'train_fuse3_loss': trainer.log_loss['fuse3_loss'],
                        'train_fuse2_loss': trainer.log_loss['fuse2_loss'],
                        'train_fuse1_loss': trainer.log_loss['fuse1_loss'],
                    },idx)

                # 定期记录训练准确率（默认每40步）
                if idx % cfg.vis_train_acc_every == 0:
                    trainer.acc_op(pred[0], target)
                    trainer.vis.log(trainer.log_acc, 'train_acc')
                    trainer.vis.plot_many({
                        'train_mask_acc': trainer.log_acc['mask_acc'],
                        'train_mask_pos_acc': trainer.log_acc['mask_pos_acc'],
                        'train_mask_neg_acc': trainer.log_acc['mask_neg_acc'],
                    },idx)

                # 定期可视化训练样本（默认每120步）
                if idx % cfg.vis_train_img_every == 0:
                    trainer.vis.img_many({
                        'train_img_big': data.cpu(),
                        # 'train_output': torch.sigmoid(pred[0].contiguous().cpu()),
                        # 'train_lab_big': target.unsqueeze(1).cpu(),
                        # 'train_fuse5': torch.sigmoid(pred[1].contiguous().cpu()),
                        # 'train_fuse4': torch.sigmoid(pred[2].contiguous().cpu()),
                        # 'train_fuse3': torch.sigmoid(pred[3].contiguous().cpu()),
                        # 'train_fuse2': torch.sigmoid(pred[4].contiguous().cpu()),
                        # 'train_fuse1': torch.sigmoid(pred[5].contiguous().cpu()),
                    },idx)

                # --------------------- 定期验证 ------------------- #
                if idx % cfg.val_every == 0:
                    trainer.vis.log('Start Val %d ....' % idx, 'train info')
                    model.eval()  # 切换评估模式

                    # 初始化验证指标
                    val_loss = {
                        'eval_total_loss': 0,
                        'eval_output_loss': 0,
                        'eval_fuse5_loss': 0,
                        'eval_fuse4_loss': 0,
                        'eval_fuse3_loss': 0,
                        'eval_fuse2_loss': 0,
                        'eval_fuse1_loss': 0,
                    }
                    val_acc = {
                        'mask_acc': 0,
                        'mask_pos_acc': 0,
                        'mask_neg_acc': 0,
                    }

                    bar.set_description('Epoch %d --- Evaluation --- :' % epoch)

                    with torch.no_grad():  # 禁用梯度计算
                        for idx, (img, lab) in enumerate(val_loader, start=1):
                            val_data = img.type(torch.cuda.FloatTensor).cuda()
                            val_target = lab.type(torch.cuda.FloatTensor).cuda()

                            # 验证推理
                            val_pred = trainer.val_op(val_data, val_target)
                            trainer.acc_op(val_pred[0], val_target)

                            # 累计验证指标
                            val_loss['eval_total_loss'] += trainer.log_loss['total_loss']
                            val_loss['eval_output_loss'] += trainer.log_loss['output_loss']
                            val_loss['eval_fuse5_loss'] += trainer.log_loss['fuse5_loss']
                            val_loss['eval_fuse4_loss'] += trainer.log_loss['fuse4_loss']
                            val_loss['eval_fuse3_loss'] += trainer.log_loss['fuse3_loss']
                            val_loss['eval_fuse2_loss'] += trainer.log_loss['fuse2_loss']
                            val_loss['eval_fuse1_loss'] += trainer.log_loss['fuse1_loss']

                            val_acc['mask_acc'] += trainer.log_acc['mask_acc']
                            val_acc['mask_pos_acc'] += trainer.log_acc['mask_pos_acc']
                            val_acc['mask_neg_acc'] += trainer.log_acc['mask_neg_acc']

                        # 可视化验证结果
                        else:
                            # 可视化验证样本
                            trainer.vis.img_many({
                                'eval_img': val_data.cpu(),
                                # 'eval_output': torch.sigmoid(val_pred[0].contiguous().cpu()),
                                # 'eval_lab': val_target.unsqueeze(1).cpu(),
                                # 'eval_fuse5': torch.sigmoid(val_pred[1].contiguous().cpu()),
                                # 'eval_fuse4': torch.sigmoid(val_pred[2].contiguous().cpu()),
                                # 'eval_fuse3': torch.sigmoid(val_pred[3].contiguous().cpu()),
                                # 'eval_fuse2': torch.sigmoid(val_pred[4].contiguous().cpu()),
                                # 'eval_fuse1': torch.sigmoid(val_pred[5].contiguous().cpu()),
                            },idx)

                            # 绘制验证损失曲线
                            trainer.vis.plot_many({
                                'eval_total_loss': val_loss['eval_total_loss'] / idx,
                                'eval_output_loss': val_loss['eval_output_loss'] / idx,
                                'eval_fuse5_loss': val_loss['eval_fuse5_loss'] / idx,
                                'eval_fuse4_loss': val_loss['eval_fuse4_loss'] / idx,
                                'eval_fuse3_loss': val_loss['eval_fuse3_loss'] / idx,
                                'eval_fuse2_loss': val_loss['eval_fuse2_loss'] / idx,
                                'eval_fuse1_loss': val_loss['eval_fuse1_loss'] / idx,
                            },idx)

                            # 绘制验证准确率曲线
                            trainer.vis.plot_many({
                                'eval_mask_acc': val_acc['mask_acc'] / idx,
                                'eval_mask_neg_acc': val_acc['mask_neg_acc'] / idx,
                                'eval_mask_pos_acc': val_acc['mask_pos_acc'] / idx,
                            },idx)

                            # 根据验证指标保存最佳模型
                            if (cfg.save_pos_acc < (val_acc['mask_pos_acc'] / idx) and
                                    cfg.save_acc < (val_acc['mask_acc'] / idx)):
                                cfg.save_pos_acc = val_acc['mask_pos_acc'] / idx
                                cfg.save_acc = val_acc['mask_acc'] / idx
                                trainer.saver.save(
                                    model,
                                    tag='%s_epoch(%d)_acc(%0.5f_%0.5f)' % (
                                        cfg.name, epoch, cfg.save_pos_acc, cfg.save_acc)
                                )
                                trainer.vis.log(
                                    'Save Model %s_epoch(%d)_acc(%0.5f_%0.5f)' % (
                                        cfg.name, epoch, cfg.save_pos_acc, cfg.save_acc),
                                    'train info'
                                )

                    bar.set_description('Epoch %d --- Training --- :' % epoch)
                    model.train()  # 恢复训练模式

            # 每epoch保存检查点
            if epoch != 0:
                trainer.saver.save(model, tag='%s_epoch(%d)' % (cfg.name, epoch))
                trainer.vis.log('Save Model -%s_epoch(%d)' % (cfg.name, epoch), 'train info')

    # ==================== 异常处理阶段 ==================== #
    except KeyboardInterrupt:
        # 键盘中断时自动保存当前模型
        trainer.saver.save(model, tag='Auto_Save_Model')
        print('\n Catch KeyboardInterrupt, Auto Save final model : %s' % trainer.saver.show_save_pth_name)
        trainer.vis.log(
            'Catch KeyboardInterrupt, Auto Save final model : %s' % trainer.saver.show_save_pth_name,
            'train info'
        )
        trainer.vis.log('Training End!!')
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)

if __name__ == '__main__':
    main()
